package s3

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/s3/types"
	"github.com/pkg/errors"
	"github.com/rs/zerolog/log"
	"github.com/sourcegraph/conc/pool"
	"io"
	"net/http"
	"oss-bench/pkg/memfile"
	"path"
	"strconv"
	"time"
)

type Storage struct {
	*AccessOptions
	*BenchOptions
	s3client *s3.Client
}

func NewS3Storage(ao *AccessOptions, bo *BenchOptions) *Storage {
	ao.Fill()
	bo.Fill()
	s3s := new(Storage)
	s3s.AccessOptions = ao
	s3s.BenchOptions = bo
	credential := credentials.NewStaticCredentialsProvider(s3s.AccessKey, s3s.SecretKey, "")
	cfg := aws.NewConfig()
	cfg.RetryMaxAttempts = 2
	cfg.Credentials = credential
	cfg.Region = s3s.Region
	cfg.EndpointResolverWithOptions = aws.EndpointResolverWithOptionsFunc(func(service, region string, options ...interface{}) (aws.Endpoint, error) {
		return aws.Endpoint{
			URL:               s3s.EndpointUrl,
			HostnameImmutable: true,
			SigningRegion:     s3s.Region,
		}, nil
	})
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	cfg.HTTPClient = &http.Client{Transport: tr}
	s3s.s3client = s3.NewFromConfig(*cfg)
	return s3s
}

func (s *Storage) Check(ctx context.Context) error {
	output, err := s.s3client.GetBucketVersioning(ctx, &s3.GetBucketVersioningInput{Bucket: &s.Bucket})
	if err != nil {
		log.Err(err).Interface("out", output).Send()
		return err
	}
	return nil
}

func (s *Storage) PutData(ctx context.Context) error {
	s0 := time.Now()
	p := pool.New().WithMaxGoroutines(s.PutConcurrent).WithContext(ctx).WithFirstError().WithCancelOnError()
	for i := 0; i < s.PutNum; i++ {
		name := strconv.Itoa(i)
		p.Go(func(ctx context.Context) error {
			return s.putOne(ctx, name)
		})
	}
	err := p.Wait()
	if err != nil {
		log.Err(err).Send()
		return err
	}
	log.Info().Int("put-concurrent", s.PutConcurrent).TimeDiff("duration-ms", time.Now(), s0).Msg("put data")
	return nil
}

func (s *Storage) putOne(ctx context.Context, name string) error {
	data := s.buildData()
	f := memfile.NewFile(name)
	f.Write(data)
	key := path.Join(s.PutPrefix, name)
	err := s.Store(ctx, f, key)
	if err != nil {
		return err
	}
	return nil
}

func (s *Storage) buildData() []byte {
	data := make([]byte, 0, s.PutSize)
	for i := 0; i < s.PutSize; i++ {
		data = append(data, 'a')
	}
	return data
}

func (s *Storage) Store(ctx context.Context, f io.Reader, key string) error {
	params := &s3.PutObjectInput{
		Bucket: &s.Bucket,
		Key:    &key,
		Body:   f,
	}
	_, err := s.s3client.PutObject(ctx, params)
	if err != nil {
		return errors.Wrap(err, "")
	}
	return nil
}

func (s *Storage) Load(ctx context.Context, key string) (io.ReadCloser, error) {
	params := &s3.GetObjectInput{
		Bucket: &s.Bucket,
		Key:    &key,
	}
	object, err := s.s3client.GetObject(ctx, params)
	if err != nil {
		return nil, err
	}
	return object.Body, nil
}

func (s *Storage) Remove(ctx context.Context, key string) error {
	if key == "" {
		return nil
	}
	params := &s3.DeleteObjectInput{
		Bucket: &s.Bucket,
		Key:    &key,
	}
	_, err := s.s3client.DeleteObject(ctx, params)
	if err != nil {
		return err
	}
	return nil
}

func (s *Storage) getObjectIds(ctx context.Context, prefix string) ([]types.ObjectIdentifier, error) {
	keys := make([]types.ObjectIdentifier, 0, s.PutNum)
	params := &s3.ListObjectsV2Input{
		Bucket:            &s.Bucket,
		ContinuationToken: nil,
		Prefix:            &prefix,
		StartAfter:        nil,
	}
	next := true
	for next {
		listObjects, err := s.s3client.ListObjectsV2(ctx, params)
		if err != nil {
			return nil, err
		}
		for _, obj := range listObjects.Contents {
			keys = append(keys, types.ObjectIdentifier{
				Key: obj.Key,
			})
		}
		next = listObjects.IsTruncated
		params.ContinuationToken = listObjects.NextContinuationToken
	}
	return keys, nil
}

func (s *Storage) getObjectIdsWithRemove(ctx context.Context, prefix string) error {
	params := &s3.ListObjectsV2Input{
		Bucket:            &s.Bucket,
		ContinuationToken: nil,
		Prefix:            &prefix,
		StartAfter:        nil,
	}
	next := true
	for next {
		listObjects, err := s.s3client.ListObjectsV2(ctx, params)
		if err != nil {
			return err
		}
		next = listObjects.IsTruncated
		params.ContinuationToken = listObjects.NextContinuationToken
		identifiers := make([]types.ObjectIdentifier, 0, len(listObjects.Contents))
		if listObjects.Contents == nil || len(listObjects.Contents) == 0 {
			return nil
		}
		for _, obj := range listObjects.Contents {
			identifiers = append(identifiers, types.ObjectIdentifier{Key: obj.Key})
		}
		err = s.removeAll(ctx, identifiers)
		if err != nil {
			return err
		}
	}
	return nil
}

// RemoveMany remove data by prefix
func (s *Storage) RemoveMany(ctx context.Context) error {
	err := s.checkDeleteObjectAuth(ctx, s.DeletePrefix)
	if err != nil {
		return err
	}
	s0 := time.Now()
	if s.DeleteFollow {
		err = s.getObjectIdsWithRemove(ctx, s.DeletePrefix)
		log.Info().TimeDiff("duration-ms", time.Now(), s0).Msg("delete data")
		if err != nil {
			return err
		}
		return nil
	}
	identifiers, err := s.getObjectIds(ctx, s.DeletePrefix)
	if err != nil {
		return err
	}
	s1 := time.Now()
	log.Info().TimeDiff("duration-ms", s1, s0).Msg("get delete objects")
	if len(identifiers) == 0 {
		return nil
	}
	log.Info().Int("len", len(identifiers)).Msg("identifiers")
	err = s.removeBatch(ctx, identifiers, s.DeleteBatchSize)
	if err != nil {
		return err
	}
	s2 := time.Now()
	log.Info().Int("delete-concurrent", s.DeleteConcurrent).TimeDiff("delete-duration-ms", s2, s1).TimeDiff("total-duration-ms", s2, s0).Msg("delete data")
	return nil
}

// removeBatch split by batchSize to remove identifiers.
func (s *Storage) removeBatch(ctx context.Context, identifiers []types.ObjectIdentifier, batchSize int) error {
	p := pool.New().WithContext(ctx).WithMaxGoroutines(s.DeleteConcurrent)
	epoch := len(identifiers) / batchSize
	for i := 0; i <= epoch; i++ {
		start := i * batchSize
		stop := start + batchSize
		if stop > len(identifiers) {
			stop = len(identifiers)
		}
		if start == stop {
			break
		}
		p.Go(
			func(ctx context.Context) error {
				return s.removeAll(ctx, identifiers[start:stop])
			},
		)
	}
	err := p.Wait()
	if err != nil {
		return err
	}
	return nil
}

// removeAll  remove all given identifiers
func (s *Storage) removeAll(ctx context.Context, identifiers []types.ObjectIdentifier) error {
	s0 := time.Now()
	objs := new(types.Delete)
	objs.Objects = identifiers
	params := &s3.DeleteObjectsInput{
		Bucket: &s.Bucket,
		Delete: objs,
	}
	resp, err := s.s3client.DeleteObjects(ctx, params)
	if err != nil {
		return err
	}
	if len(resp.Errors) == 0 {
		return nil
	}
	err = s3ErrAggregate(resp.Errors)
	log.Debug().TimeDiff("Duration ms", time.Now(), s0).Msg("remove one epoch")
	return err
}

func (s *Storage) checkDeleteObjectAuth(ctx context.Context, key string) error {
	input := &s3.DeleteObjectInput{
		Bucket: &s.Bucket,
		Key:    &key,
	}
	_, err := s.s3client.DeleteObject(ctx, input)
	err = errors.WithMessage(err, "checkDeleteObjectAuth")
	return err
}

func s3ErrAggregate(errs []types.Error) error {
	m := make(map[string]string)
	for _, err := range errs {
		m[*err.Key] = *err.Message
	}
	raw, _ := json.Marshal(m)
	return errors.New(string(raw))
}
